// Copyright (c) 2023 Presto Labs Pte. Ltd.
// Author: ziyan

#include <cxxopts.hpp>
#include "coin2/app/driver.h"
#include "coin2/base/config/args.h"
#include "coin2/base/log.h"
#include "coin2/base/time.h"
#include "coin2/strategy/strategy.h"
#include <highfive/H5File.hpp>
#include "coin2/flow/writer.h"
#include "coin2/exchange/base/executor/topic_distributor.h"
#include "coin/proto/coin_market_query.pb.h"


using namespace coin2::exchange::base::symbology;
using namespace coin2::base::config;
using namespace coin2::exchange::feed;
using namespace coin2::app;
using namespace HighFive;
using coin2::exchange::base::feed::FeedUpdate;
using coin2::exchange::base::symbology::IProduct;
using coin2::exchange::base::market::MarketExchangeApi;
using coin2::exchange::base::executor::TopicRecordProcessor;
using coin2::exchange::base::market::ExchangeTypeToString;
using coin2::exchange::base::market::MarketTypeToString;

struct ExchangeTicker {
  std::unordered_map<std::string, int> symbol_id;
  std::vector<std::string> norm_symbol;
  std::vector<int64_t> exchange_timestamp;
  std::vector<double> mid_price;
};


class TickerProviderStrategy : public ::coin2::strategy::IStrategy {
 public:
  TickerProviderStrategy(
    const bool& write_h5,
    const bool& write_proto,
    const std::string& h5_out_dir,
    const std::string& proto_out_dir,
    const int& resolution )
    : write_h5_(write_h5),
    write_proto_(write_proto),
    h5_out_dir_(h5_out_dir),
    proto_out_dir_(proto_out_dir),
    resolution_(resolution){
  }

  void onBookFeed(const FeedUpdate& upd) override {
    auto& book = upd.book();
    auto& exchange_ticker = exchange_tickers_.at(upd.mea());
    int id = exchange_ticker.symbol_id.at(upd.product().absolute_norm());
    if (book.Ask0() && book.Bid0()) {
      exchange_ticker.mid_price[id] = (book.Ask0()->price + book.Bid0()->price) / 2;
      exchange_ticker.exchange_timestamp[id] = upd.timestamp();
    } else if (upd.mea().String() == "Options.Deribit.v2") {
      exchange_ticker.mid_price[id]=std::numeric_limits<double>::quiet_NaN();
      exchange_ticker.exchange_timestamp[id] = upd.timestamp();
    }
  }

  void Init(Driver* driver) override {
    for (auto& subsystem : driver->feed()->GetAllFeedSubsystems()) {
      MarketExchangeApi mea = subsystem->mea();
      std::unordered_map<std::string, int> symbol_id;
      std::vector<std::string> norm_symbol;
      int id = 0;
      for (auto& product : subsystem->GetProducts()){
        auto& symbol = product->absolute_norm();
        symbol_id.insert({symbol,id});
        norm_symbol.push_back(product->symbol());
        id++;
      }
      std::vector<double> mid_price(symbol_id.size());
      std::fill(mid_price.begin(), mid_price.end(), std::numeric_limits<double>::quiet_NaN());
      std::vector<int64_t> exchange_timestamp(symbol_id.size());
      ExchangeTicker exchange_ticker = ExchangeTicker{symbol_id, norm_symbol, exchange_timestamp, mid_price};
      exchange_tickers_.insert({mea, exchange_ticker});
    }
    machine_name_ = driver->config().feed().common().archive().machine();
    write_to_file_timer_ = std::make_unique<boost::asio::deadline_timer>(*driver->live_executor()->MainIoc());
    write_to_file_timer_->expires_at(next_write_time());
    write_to_file_timer_->async_wait(std::bind(&TickerProviderStrategy::write, this));
  }
  
  private:
  bool write_h5_;
  bool write_proto_;
  std::string h5_out_dir_;
  std::string proto_out_dir_;
  int resolution_;
  std::string h5_file_name_;
  std::string machine_name_;
  std::unique_ptr<boost::asio::deadline_timer> write_to_file_timer_;
  std::unordered_map<MarketExchangeApi, ExchangeTicker> exchange_tickers_;
  std::unordered_map<MarketExchangeApi, std::unique_ptr<FlowWriter>> writers_;

  void write(){
    boost::posix_time::ptime index_pt = write_to_file_timer_->expires_at();
    int64_t index_ts = TimestampFromPtime(index_pt);
    if (write_h5_){
      write_h5(index_ts);
    }
    if (write_proto_){
      GenerateTicker(index_ts);
    }
    write_to_file_timer_->expires_at(index_pt + boost::posix_time::seconds(resolution_));
    write_to_file_timer_->async_wait(std::bind(&TickerProviderStrategy::write, this));
  }

   void GenerateTicker(int64_t index_ts) {
    for (auto& [mea, ticker] : exchange_tickers_){
      int64_t timestamp = TopicRecordProcessor::GetTimestamp();
      coin::proto::SubscriptionLog proto;
      proto.set_type(coin::proto::SubscriptionLog::TICKER);
      proto.mutable_request()->set_market_type(MarketTypeToString(mea.market));
      proto.mutable_request()->set_exchange(ExchangeTypeToString(mea.exchange));
      proto.mutable_request()->set_api_version(mea.api);
      proto.mutable_ticker()->set_exchange(ExchangeTypeToString(mea.exchange));
      proto.mutable_ticker()->set_market_type(MarketTypeToString(mea.market));
      proto.mutable_ticker()->set_api_version(mea.api);
      proto.mutable_ticker()->set_query_ts(timestamp);
      for (const auto& [product, id] : ticker.symbol_id) {
        if (!std::isnan(ticker.mid_price[id])){
          auto product_ticker = proto.mutable_ticker()->add_each_ticker();
          product_ticker->set_symbol(ticker.norm_symbol[id]);
          product_ticker->set_mid(ticker.mid_price[id]);
          product_ticker->set_last(ticker.mid_price[id]);
          product_ticker->set_exchange_timestamp(ticker.exchange_timestamp[id]);
          product_ticker->set_interval_timestamp(index_ts);
        }
      }
      std::string buf;
      buf.resize(orio::record::TopicRecord::struct_size_with_data(proto.ByteSize()));
      auto* r = reinterpret_cast<orio::record::TopicRecord*>(buf.data());
      r->data_len = proto.ByteSize();
      r->queue_seq = 0;
      r->topic_id = 0;
      r->topic_seq = 0;
      r->timestamp = timestamp;

      std::string buf2;
      proto.SerializeToString(&buf2);
      const auto* data = reinterpret_cast<const std::byte*>(buf2.data());
      std::copy(data, data + r->data_len, r->data);

      if (writers_.find(mea) == std::end(writers_)) {
        writers_[mea] = std::make_unique<FlowWriter>(
            proto_out_dir_, machine_name_, mea.SnakeString());
      }

      FlowWriter* writer = writers_[mea].get();
      writer->Write(*r);
      writer->Flush();
    }
  }


  void write_h5(int64_t index_ts){
    for (auto& [mea, ticker] : exchange_tickers_){
      std::string out_path = h5_out_dir_ + "/" + mea.String() + "/" + YmdFromTimestamp(index_ts) + "/";
      std::filesystem::path fp = (out_path + h5_file_name_);
      if (h5_file_name_ != "" && std::filesystem::exists(fp)) {
        File file(fp, File::ReadWrite);
        append_mid_price(file, ticker.mid_price);
        appen_timestamp(file, index_ts);
      } else {
        LOG(INFO) << "create "<< out_path.c_str();
        std::filesystem::create_directories(out_path.c_str());
        h5_file_name_ = "realtime_mid_price_" + std::to_string(index_ts / 1'000'000'000LL) + ".h5";
        fp = (out_path + h5_file_name_);
        File file(fp, File::Create);
        init_mid_price(file, ticker.mid_price);
        init_universe(file, ticker.symbol_id);
        init_timestamp(file, index_ts);
      }
    }
  }

  void init_universe(File& file, std::unordered_map<std::string, int>& symbol_id){        
    std::vector<std::string> universe(symbol_id.size());
    for (auto& [symbol,id] : symbol_id){
      universe[id] = symbol;
    }
    file.createDataSet("universe", universe);
  }

  void init_timestamp(File& file,int64_t index_ts){        
    DataSpace timestamp_dataspace = DataSpace({1}, {DataSpace::UNLIMITED});
    DataSetCreateProps timestamp_props;
    timestamp_props.add(Chunking(std::vector<hsize_t>{1}));
    DataSet timestamp_dataset = file.createDataSet<int64_t>("timestamp", timestamp_dataspace ,timestamp_props);
    timestamp_dataset.write(index_ts);
  }

  void init_mid_price(File& file, std::vector<double>& mid_price){
    auto width = mid_price.size();
    DataSpace mid_price_close_dataspace = DataSpace({1, width}, {DataSpace::UNLIMITED, width});
    DataSetCreateProps mid_price_close_props;
    mid_price_close_props.add(Chunking(std::vector<hsize_t>{1, width}));
    DataSet mid_price_dataset = file.createDataSet<double>("close_mid", mid_price_close_dataspace, mid_price_close_props);
    mid_price_dataset.write(mid_price);
  }

  void append_mid_price(File& file, std::vector<double>& mid_price){
    auto mid_price_close_dataset = file.getDataSet("close_mid");
    std::vector<size_t> dims = mid_price_close_dataset.getDimensions();
    mid_price_close_dataset.resize({dims[0] + 1, dims[1]});
    mid_price_close_dataset.select({dims[0], 0}, {1, dims[1]}).write(mid_price);
  }

  void appen_timestamp(File& file, int64_t index_ts){
    auto ts_dataset = file.getDataSet("timestamp");
    size_t ts_size = ts_dataset.getElementCount();
    ts_dataset.resize({ts_size + 1});
    ts_dataset.select({ts_size},{1}).write(index_ts);
  }

  boost::posix_time::ptime next_write_time(){
    time_t now = time(0);
    time_t next_ts  = now - now % resolution_ + resolution_;
    boost::posix_time::ptime start_time = boost::posix_time::from_time_t(next_ts);
    return start_time;
  }
};

int main(int argc, char* argv[]) {
  google::InitGoogleLogging(argv[0]);

  DriverConfig app;

  cxxopts::Options opt("Provider", "Provide interval mid price");
  opt.add_options()("resolution","length of interval in sec",cxxopts::value<int>()->default_value("60"));
  opt.add_options()("write_h5","write h5 file",cxxopts::value<bool>()->default_value("false"));
  opt.add_options()("write_proto","write proto file",cxxopts::value<bool>()->default_value("false"));
  opt.add_options()("h5_out_dir","ticker h5 file out directory",cxxopts::value<std::string>()->default_value("/remote/iosg/data-2/buckets/feed.derived.realtime.mid_price_interval_h5/coin"));
  opt.add_options()("proto_out_dir","ticker proto file out directory",cxxopts::value<std::string>()->default_value("/home/ec2-user/data/sub_proto_log"));

  AddDriverOptions(&opt);

  auto res = opt.parse(argc, argv);
  ParseDriverOptions(res, &app);

  LOG(INFO) << app.DebugString();

  TickerProviderStrategy provider(
      res["write_h5"].as<bool>(),
      res["write_proto"].as<bool>(),
      res["h5_out_dir"].as<std::string>(),
      res["proto_out_dir"].as<std::string>(),
      res["resolution"].as<int>());
  Driver driver(app, &provider);
  driver.Run();

  return 0;
}
